# -*- coding: utf-8 -*-
# created on 2016/11/30
from mathsolver.functions.base import *
from mathsolver.functions.shulie.Fundamental import *
from mathsolver.functions.sympy_utils import default_symbol
from sympy import Poly


class ShuLieUpdate001(BaseFunction):
    """
    已知数列{a_{n}}中,a_{1}=2,{a}_{n+1}=(\\sqrt{2}-1)({a}_{n}+2),求{a_{n}}的通项公式.
    """
    def solver(self, *args):
        sl = BaseSequence(args[0].value)
        a_name = sl.aName
        sl.name = str(a_name)
        # 如果self.known里面含有该数列, 不新建数列了
        assert self.search(a_name) is None
        self.output.append(sl)
        return self


class ShuLieUpdate002(BaseFunction):
    """
    设数列满足a_{1}=2,a_{n+1}-a_{n}=3•2^{2n-1},求数列{a_{n}}的通项公式.
    """
    def solver(self, *args):
        assert len(args) == 1
        assert isinstance(args[0], (BaseSeqEqs, BaseSeqEq))
        if isinstance(args[0], BaseSeqEq):
            eq = args[0].sympify()
            eqs = [eq]
        else:
            eqs = args[0].sympify()
        names = []
        for eq in eqs:
            expr = eq[0] - eq[1]
            p = r'[a-zA-Z]_\(.*?\)'
            m = re.findall(p, str(expr))
            m = set(m)
            m = list(m)
            names.extend(m)
        for ti, tv in enumerate(names):
            tv = sympify(tv)
            names[ti] = str(tv.func)
        names = list(set(names))
        a_name = ""
        s_name = ""
        if len(names) == 1:
            name = names[0]
            p = r'[a-z]_'
            m = re.findall(p, str(name))
            m = set(m)
            m = list(m)
            if m:
                m = m[0]
                a_name = m
            else:
                p = r'[A-Z]_'
                m = re.findall(p, str(name))
                m = set(m)
                m = list(m)
                if m:
                    m = m[0]
                    s_name = m
        elif len(names) == 2:
            for name in names:
                p = r'[a-z]_'
                m = re.findall(p, str(name))
                m = set(m)
                m = list(m)
                if m:
                    m = m[0]
                    a_name = m
                else:
                    p = r'[A-Z]_'
                    m = re.findall(p, str(name))
                    m = set(m)
                    m = list(m)
                    if m:
                        m = m[0]
                        s_name = m
        else:
            raise Exception("length of names is wrong")
        name = a_name + s_name
        if a_name != "":
            sl = BaseSequence(a_name + "(n)")
            sl.name = name
            self.output.append(sl)
        else:
            raise Exception("to do")
        return self


class ShuLieUpdate(BaseFunction):
    CLS = [ShuLieUpdate001, ShuLieUpdate002]

    def solver(self, *args):
        r = None
        for cl in ShuLieUpdate.CLS:
            try:
                r = cl(known=self.known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r


def sl_sn_max_min_update(seqs, ineq, known):
    from mathsolver.functions.shulie.basic import sl_search_known
    left, op, right = ineq
    expr = left - right
    expr_xiangs = sl_xiangs001(expr)
    flag = False
    s_name = None
    answer_index = None
    answer = None
    if len(expr_xiangs) == 1:
        new_item = sympify(expr_xiangs[0])
        new_item_name = new_item.func
        new_item_index = new_item.args[0]
        sl = sl_search_known(known, str(new_item_name))
        s_name = sl.sName
        n_name = sl.nName
        if str(new_item_name) == str(s_name) and str(new_item_index) == str(n_name):
            answers = solve(expr, new_item)
            if len(answers) == 1:
                answer = answers[0]
                minus = expr - new_item + answer
                if minus == 0:
                    answer_name = answer.func
                    answer_index = answer.args[0]
                    if str(answer_name) == str(s_name):
                        if str(answer_index).find(str(n_name)) < 0:
                            if sl.type == "DengCha":
                                flag = True
    assert s_name is not None
    if flag:
        seqs.append((s_name(answer_index - 1), op, answer))
        seqs.append((s_name(answer_index + 1), op, answer))
    else:
        seqs.append(ineq)
    return seqs


class ShuLieEqUpdate(BaseFunction):
    """
    记等差数列{a_{n}}的前n项和为S_{n},设S_{3}=12,2a_{1},a_{2},a_{3}+1成等比数列,求S_{n}.
    """
    def solver(self, *args):
        eq = args[0].sympify()
        if isinstance(args[0], (BaseSeqEq, BaseEq, BaseSeqIneq, BaseIneq)):
            answer = {}
            if 'Sequations' in self.known:
                seqs = self.search('Sequations')
            else:
                seqs = []
            if isinstance(args[0], (BaseSeqIneq, BaseIneq)):
                seqs = sl_sn_max_min_update(seqs, eq, self.known)
            else:
                seqs.append(eq)
            answer['Sequations'] = seqs
            self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieEqsUpdate(BaseFunction):
    """
    等差数列{a_{n}}中,a_{5}+a_{7}=16,a_{3}=4,则a_{9}=.
    """
    def solver(self, *args):
        eqs = args[0].sympify()
        if isinstance(args[0], (BaseEqs, BaseSeqEqs, BaseMultiple, BaseIneqs, BaseSeqIneqs)):
            answer = {}
            if 'Sequations' in self.known:
                seqs = self.search('Sequations')
            else:
                seqs = []
            for eq in eqs:
                if len(eq) == 2:
                    seqs.append(eq)
                elif len(eq) == 3:
                    if isinstance(eq[1], str):
                        seqs.append(eq)
                    else:
                        seqs.append([eq[0], eq[2]])
                        seqs.append([eq[1], eq[2]])
            answer['Sequations'] = seqs
            self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGenEqUpdate(BaseFunction):
    def solver(self, *args):
        eq = args[0].sympify()
        answer = args[1].sympify()
        symbol = default_symbol(eq[0] - eq[1])
        eq = [eq[0].subs({symbol: answer}), eq[1].subs({symbol: answer})]
        if isinstance(args[0], (BaseSeqEq, BaseEq)):
            answer = {}
            if 'Sequations' in self.known:
                seqs = self.search('Sequations')
            else:
                seqs = []
            seqs.append([eq[0], eq[1]])
            answer['Sequations'] = seqs
            self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieTongXiangUpdate(BaseFunction):
    """
    已知数列的通项a_{n}=-5n+2,则其前n项和S_{n}=.
    """
    def solver(self, *args):
        eq = args[0].sympify()
        fn = (eq[0] - eq[1]).expand().simplify()
        fn = sl_new_an(fn)
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(fn))
        if m:
            if len(m) == 1:
                item = m[0]
                new_item = sympify(item)
                new_item_name = new_item.func
                answer = solve(fn, new_item)
                new_item_value = answer[0]
                sl = self.search(str(new_item_name))
                sl.an.append(new_item_value)
        return self


class ShuLieAnIneqUpdate(BaseFunction):
    """
    正项数列{a_{n}}满足:a_{n}^2-(2n - 1)a_{n}-2n=0.求数列{a_{n}}的通项公式.
    """
    def solver(self, *args):
        assert len(args) == 2
        seq = args[0].sympify()
        text = args[1]
        answer = {}
        if 'Sequations' in self.known:
            seqs = self.search('Sequations')
        else:
            seqs = []
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(seq))
        m = set(m)
        m = list(m)
        assert len(m) == 1
        new_item = sympify(m[0])
        new_item_name = new_item.func
        sl = self.search(str(new_item_name))
        a_name = sl.aName
        n_name = sl.nName
        if text.find("正") >= 0:
            a_name = sympify(a_name)
            seqs.append([a_name(n_name), ">", S.Zero])
            answer['Sequations'] = seqs
            self.output.append(BaseShuLieValue(answer))
            self.steps.append(["", "%s" % BaseIneq([a_name(n_name), ">", S.Zero]).printing()])
        else:
            raise Exception("to do")
        return self


class ShuLieEqInfUpdate(BaseFunction):
    """
    已知等比数列{a_{n}}中,a_{1}=\\frac{1}{3},公比q=\\frac{1}{3}.设b_{n}=log_{3}a_{1}+log_{3}a_{2}+…+log_{3}a_{n},求数列{b_{n}}的通项公式.
    """
    def solver(self, *args):
        assert len(args) == 1
        assert isinstance(args[0], BaseSeqEqInf)
        left, right = args[0].value
        try:
            left = sympify(left)
        except Exception:
            left = left
        try:
            right = sympify(right)
        except Exception:
            right = right
        answer = {}
        if 'InfSequations' in self.known:
            seqs = self.search('InfSequations')
        else:
            seqs = []
        seqs.append([left, right])
        answer['InfSequations'] = seqs
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieAnIntegerUpdate(BaseFunction):
    """
    等差数列{a_{n}}的前n项和为S_{n}.已知a_{1}=10,a_{2}为整数,S_{n}≤S_{4}.求{a_{n}}的通项公式.
    """
    def solver(self, *args):
        assert len(args) == 1
        seq = args[0].sympify()
        answer = {}
        if 'Sequations' in self.known:
            seqs = self.search('Sequations')
        else:
            seqs = []
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(seq))
        m = set(m)
        m = list(m)
        assert len(m) == 1
        new_item = sympify(m[0])
        new_item_name = new_item.func
        new_item_index = new_item.args[0]
        sl = self.search(str(new_item_name))
        sl_type = sl.type
        a_name = sl.aName
        a_name = sympify(a_name)
        if sl_type == str("DengCha"):
            d_value = sl.dValue
            if str(new_item_name) == str(a_name):
                if d_value is None:
                    d_value = 'd'
                    d_value = sympify(d_value)
                new_item_value = a_name(1) + (new_item_index - 1) * d_value
                seq = seq.subs({new_item: new_item_value})
        seqs.append(BaseBelong([seq, "\\in", "Z"]))
        answer['Sequations'] = seqs
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieWeiDaDingLiUpdate(BaseFunction):
    """
    在等比数列{a_{n}}中,若a_{3},a_{9}是方程3x^{2}-11x+9=0的两根,则a_{6}的值是()
    """
    def solver(self, *args):
        x1, x2 = args[0].sympify()
        f1, f2 = args[1].sympify()
        expr = (f1 - f2).expand().simplify()
        symbol = default_symbol(expr)
        poly = Poly(expr, symbol)
        a, b, c = poly.all_coeffs()
        self.steps.append(["", "由韦达定理，得"])
        eq1 = [x1 + x2, -b / a]
        eq2 = [x1 * x2, c / a]
        eqs = [eq1, eq2]
        self.steps.append(["", "%s" % BaseEqs(eqs).printing()])
        answer = {}
        if 'Sequations' in self.known:
            seqs = self.search('Sequations')
        else:
            seqs = []
        seqs.append(eq1)
        seqs.append(eq2)
        answer['Sequations'] = seqs
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieSnEqualSnUpdate(BaseFunction):
    """
    等差数列\\{a_{n}\\}前9项的和等于前4项的和.若a_{1}=1,a_{k}+a_{4}=0,则k=.
    """
    def solver(self, *args):
        seq = args[0].sympify()
        num1 = args[1].sympify()
        num2 = args[2].sympify()
        answer = {}
        if 'Sequations' in self.known:
            seqs = self.search('Sequations')
        else:
            seqs = []
        seq_name = seq.func
        sl = self.search(str(seq_name))
        s_name = sl.sName
        if s_name is None:
            s_name = "S_"
            del self.known[str(seq_name)]
            sl.sName = s_name
            sl.name = str(seq_name) + s_name
            self.output.append(sl)
        s_name = sympify(s_name)
        eq = [s_name(num1), s_name(num2)]
        seqs.append(eq)
        answer['Sequations'] = seqs
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieAnValueUpdate(BaseFunction):
    """
    等比数列3,6,12,24的公比q=.
    """
    def solver(self, *args):
        if len(args) == 1:
            values = args[0].sympify()
            sls = [self.known[item] for item in self.known if isinstance(self.known[item], BaseSequence)]
            assert len(sls) == 1
            sl = sls[0]
        else:
            sl = args[0]
            values = args[1].sympify()
        a_name = sl.aName
        a_name = sympify(a_name)
        answer = {}
        if 'Sequations' in self.known:
            seqs = self.search('Sequations')
        else:
            seqs = []
        for i, v in enumerate(values):
            seqs.append([a_name(i+1), v])
        answer['Sequations'] = seqs
        self.output.append(BaseShuLieValue(answer))
        return self
